import os
import pathlib
import random
import string

import fitz
import ocrmypdf
from PIL import Image
from loguru import logger
from pdfminer.high_level import extract_text

special_char_map = {
    0xA728: "TZ",
    0xA729: "tz",
    0xA732: "AA",
    0xA733: "aa",
    0xA734: "AO",
    0xA735: "ao",
    0xA736: "AU",
    0xA737: "au",
    0xA738: "AV",
    0xA739: "av",
    0xA73A: "AV",
    0xA73B: "av",
    0xA73C: "AY",
    0xA73D: "ay",
    0xA74E: "OO",
    0xA74F: "oo",
    0xA760: "VY",
    0xA761: "vy",
    0xFB00: "ff",
    0xFB01: "fi",
    0xFB02: "fl",
    0xFB03: "ffi",
    0xFB04: "ffl",
    0xFB05: "ft",
    0xFB06: "st"
}


async def extract_pdf_text_speed(pdf_path: str, output_dir: str, word_count_min: int, max_page: int,
                           parse_method: str) -> str:
    pdf_path = await truncate_pdf_over_pages(pdf_path, output_dir, max_page)

    if parse_method == 'auto' or parse_method == 'ocr':
        image_count = await pdf_image_count(pdf_path)

        if image_count > 0:
            output_pdf_path = os.path.join(output_dir, generate_random_string(8) + '.pdf')
            output_sidecar_path = os.path.join(output_dir, generate_random_string(8) + '_sidecar.txt')

            if parse_method == 'ocr':
                ocrmypdf.ocr(pdf_path, output_pdf_path, language=['chi_sim', 'eng'], force_ocr=True, deskew=True,
                             tesseract_config='tesseract.cfg', sidecar=output_sidecar_path)
            else:
                ocrmypdf.ocr(pdf_path, output_pdf_path, language=['chi_sim', 'eng'], redo_ocr=True,
                             tesseract_config='tesseract.cfg', sidecar=output_sidecar_path)

            if os.path.exists(output_pdf_path):
                pdf_path = output_pdf_path

    extracted_text = extract_text(pdf_path).strip()
    if len(extracted_text) < word_count_min:
        return ''
    processed_text = replace_spec_chars(extracted_text)
    text_path = os.path.join(output_dir, generate_random_string(8) + '.txt')
    with open(text_path, 'w') as f:
        f.write(processed_text)
    logger.info(f'extract_pdf_text_speed({pdf_path} -> {text_path}) success')
    return text_path


async def truncate_pdf_over_pages(input_pdf_path: str, output_dir: str, max_page: int) -> str:
    with fitz.open(input_pdf_path) as doc:
        total_pages = len(doc)
        if total_pages <= max_page:
            return input_pdf_path

        with fitz.open() as new_doc:
            for page_num in range(min(max_page, total_pages)):
                new_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)
            output_pdf_path = str(
                pathlib.Path(output_dir).joinpath(pathlib.Path(input_pdf_path).stem + "_truncated.pdf"))
            new_doc.save(output_pdf_path)
            print(f"Truncated document to {max_page} pages and saved as {output_pdf_path}.")
            return output_pdf_path


async def pdf_to_image(pdf_path: str, output_dir: str, max_page: int, target_size: int) -> str:
    base_name = os.path.splitext(os.path.basename(pdf_path))[0]
    png_output_path = os.path.join(output_dir, f"{base_name}.png")
    jpeg_output_path = os.path.join(output_dir, f"{base_name}.jpg")

    images = []
    total_height = 0
    max_width = 0

    document = fitz.open(pdf_path)
    page_len = min(len(document), max_page)

    for page_num in range(page_len):
        page = document.load_page(page_num)
        pix = page.get_pixmap()
        img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
        images.append(img)

        total_height += img.size[1]
        max_width = max(max_width, img.size[0])

    document.close()
    merged_image = Image.new('RGB', (max_width, total_height), color='white')
    current_height = 0

    for img in images:
        # 如果页面宽度小于最大宽度，则需要创建一个白色背景的图像来填充
        if img.size[0] < max_width:
            new_img = Image.new('RGB', (max_width, img.size[1]), color='white')
            new_img.paste(img, (0, 0))
            img = new_img

        merged_image.paste(img, (0, current_height))
        current_height += img.size[1]
        img.close()

    # 灰度化
    with merged_image.convert('L') as gray_img:
        merged_image.close()
        gray_img.save(png_output_path, "PNG")
        pic_size = os.path.getsize(png_output_path)
        if pic_size <= target_size:
            logger.info(f'pdf_to_image({pdf_path} -> {png_output_path}) success')
            return png_output_path

        for i in range(1, 19):
            quality = (10 - i) * 10 if i < 10 else 19 - i
            gray_img.save(png_output_path, "PNG", quality=quality)
            pic_size = os.path.getsize(png_output_path)
            if pic_size <= target_size:
                logger.info(f'pdf_to_image({pdf_path} -> {png_output_path}) success')
                return png_output_path

        os.remove(png_output_path)

        # 如果PNG失败，尝试保存为JPEG
        for i in range(1, 19):
            quality = (10 - i) * 10 if i < 10 else 19 - i
            gray_img.save(jpeg_output_path, "JPEG", quality=quality)
            pic_size = os.path.getsize(jpeg_output_path)
            if pic_size <= target_size:
                logger.info(f'pdf_to_image({pdf_path} -> {jpeg_output_path}) success')
                return jpeg_output_path

    return ''


async def pdf_to_pngs(pdf_path: str, output_dir: str) -> list[str]:
    document = fitz.open(pdf_path)
    pp = pathlib.Path(pdf_path)
    pngs = []

    for page_num in range(len(document)):
        page = document.load_page(page_num)
        pix = page.get_pixmap()
        output_image_path = f"{output_dir}/{pp.stem}_{page_num + 1}.png"
        pix.save(output_image_path)
        pngs.append(output_image_path)

    return pngs


async def pdf_image_count(pdf_path: str) -> int:
    with fitz.open(pdf_path) as document:
        total_images = 0

        for page_num in range(len(document)):
            page = document.load_page(page_num)
            images = page.get_images(full=True)

            if images:
                total_images += len(images)

        logger.info(f'pdf_image_count({pdf_path} -> {total_images}) success')
        return total_images


def generate_random_string(length: int):
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))


def replace_spec_chars(raw) -> str:
    if not raw:
        return raw

    result = []

    for c in raw:
        r = special_char_map.get(ord(c))
        if r is None:
            result.append(c)
        else:
            result.append(r)

    return ''.join(result)
